#!/usr/bin/env python3
# Copyright 2016 Open Source Robotics Foundation, Inc.
# Modified for Unitree Go2 robot laser following

import rclpy
import numpy as np
import threading
import time
from rclpy.node import Node
from rclpy.qos import QoSProfile
from sensor_msgs.msg import LaserScan
from geometry_msgs.msg import Twist
from std_msgs.msg import String, Int8

class LaserTracker(Node):
    def __init__(self):
        super().__init__('laser_tracker')
        self.last_scan = None
        qos = QoSProfile(depth=10)
        
        # Parameters
        self.declare_parameter('win_size', 2)
        self.declare_parameter('delta_dist', 0.2)
        
        # Publishers
        self.position_publisher = self.create_publisher(Twist, 'object_tracker/current_position', qos)
        self.info_publisher = self.create_publisher(String, 'object_tracker/info', qos)
        
        # Subscriber
        self.scan_subscriber = self.create_subscription(
            LaserScan,
            '/scan',
            self.register_scan,
            qos)
        
    def register_scan(self, scan_data):
        ranges = np.array(scan_data.ranges)
        sorted_indices = np.argsort(ranges)
        min_distance_id = None
        min_distance = float('inf')
        
        if self.last_scan is not None:
            for i in sorted_indices:
                temp_min_distance = ranges[i]
                a = i - 2
                b = i + 3
                window_index = np.clip([a, b], 0, len(self.last_scan))
                window = self.last_scan[window_index[0]:window_index[1]]
                
                with np.errstate(invalid='ignore'):
                    if np.any(abs(window - temp_min_distance) <= 0.2):
                        min_distance_id = i
                        min_distance = ranges[min_distance_id]
                        break
        
        self.last_scan = ranges

        if min_distance > scan_data.range_max:
            msg = String()
            msg.data = 'laser:nothing found'
            self.get_logger().warn('No object found in laser scan')
            self.info_publisher.publish(msg)
        else:
            msg = Twist()
            min_distance_angle = scan_data.angle_min + min_distance_id * scan_data.angle_increment
            
            # For Go2 robot, we'll use linear.x for distance and angular.z for angle
            msg.linear.x = float(min_distance)
            msg.angular.z = float(min_distance_angle)
            self.position_publisher.publish(msg)

class LaserFollower(Node):
    def __init__(self):
        super().__init__('laser_follower')
        
        # Parameters
        self.declare_parameter('target_distance', 1.0)  # 1 meter target distance
        self.declare_parameter('max_speed', 0.5)        # Max linear/angular speed
        
        # PID parameters
        self.declare_parameter('P_linear', 0.5)
        self.declare_parameter('I_linear', 0.0)
        self.declare_parameter('D_linear', 0.1)
        self.declare_parameter('P_angular', 1.5)
        self.declare_parameter('I_angular', 0.0)
        self.declare_parameter('D_angular', 0.2)
        
        # Get parameters
        self.target_distance = self.get_parameter('target_distance').value
        self.max_speed = self.get_parameter('max_speed').value
        
        # PID controllers
        self.distance_pid = PIDController(
            self.get_parameter('P_linear').value,
            self.get_parameter('I_linear').value,
            self.get_parameter('D_linear').value
        )
        self.angle_pid = PIDController(
            self.get_parameter('P_angular').value,
            self.get_parameter('I_angular').value,
            self.get_parameter('D_angular').value
        )
        
        # Publishers
        self.cmd_vel_publisher = self.create_publisher(Twist, 'cmd_vel', QoSProfile(depth=10))
        self.follow_flag_publisher = self.create_publisher(Int8, '/laser_follow_flag', QoSProfile(depth=10))
        
        # Subscribers
        self.position_subscriber = self.create_subscription(
            Twist,
            '/object_tracker/current_position',
            self.position_update_callback,
            QoSProfile(depth=10))
        
        self.tracker_info_subscriber = self.create_subscription(
            String,
            '/object_tracker/info',
            self.tracker_info_callback,
            QoSProfile(depth=10))
        
        # State
        self.active = True
        self.last_update_time = time.time()
        
    def position_update_callback(self, position):
        current_distance = position.linear.x
        current_angle = position.angular.z
        
        # Update PID controllers
        linear_speed = self.distance_pid.update(current_distance, self.target_distance)
        angular_speed = self.angle_pid.update(current_angle, 0.0)  # Target angle is 0 (straight ahead)
        
        # Create and publish Twist message
        cmd_vel = Twist()
        
        # Apply limits
        linear_speed = np.clip(linear_speed, -self.max_speed, self.max_speed)
        angular_speed = np.clip(angular_speed, -self.max_speed, self.max_speed)
        
        # Only move forward if object is at least 0.5m away (to avoid collisions)
        if current_distance > 0.5:
            cmd_vel.linear.x = linear_speed
            cmd_vel.angular.z = angular_speed
        else:
            cmd_vel.linear.x = 0.0
            cmd_vel.angular.z = 0.0
        
        self.cmd_vel_publisher.publish(cmd_vel)
        
        # Publish follow flag
        flag_msg = Int8()
        flag_msg.data = 1
        self.follow_flag_publisher.publish(flag_msg)
        
        self.last_update_time = time.time()
    
    def tracker_info_callback(self, info):
        if 'nothing found' in info.data:
            # Stop if no object detected
            cmd_vel = Twist()
            self.cmd_vel_publisher.publish(cmd_vel)
            self.get_logger().warn('No object detected - stopping')

class PIDController:
    def __init__(self, kp, ki, kd):
        self.kp = kp
        self.ki = ki
        self.kd = kd
        self.last_error = 0
        self.integral = 0
        self.last_time = time.time()
    
    def update(self, current_value, target_value):
        current_time = time.time()
        dt = current_time - self.last_time
        if dt <= 0:
            return 0.0
        
        error = target_value - current_value
        self.integral += error * dt
        derivative = (error - self.last_error) / dt
        
        output = self.kp * error + self.ki * self.integral + self.kd * derivative
        
        self.last_error = error
        self.last_time = current_time
        
        return output

def main(args=None):
    rclpy.init(args=args)
    
    # Create nodes
    tracker = LaserTracker()
    follower = LaserFollower()
    
    # Use MultiThreadedExecutor to run both nodes
    executor = rclpy.executors.MultiThreadedExecutor()
    executor.add_node(tracker)
    executor.add_node(follower)
    
    try:
        executor.spin()
    finally:
        tracker.destroy_node()
        follower.destroy_node()
        rclpy.shutdown()

if __name__ == '__main__':
    main()